from pydantic import BaseModel, Field
from typing import Optional, List
from datetime import datetime
from enum import Enum


class AlgorithmBase(BaseModel):
    """算法基础模式"""
    name: str
    description: Optional[str] = None


class AlgorithmCreate(AlgorithmBase):
    """算法创建模式"""
    pass


class Algorithm(AlgorithmBase):
    """算法响应模式"""
    id: int
    created_at: datetime
    
    class Config:
        orm_mode = True


class DatasetBase(BaseModel):
    """数据集基础模式"""
    name: str


class DatasetCreate(DatasetBase):
    """数据集创建模式"""
    file_path: str
    record_count: int = 0
    feature_count: int = 0
    class_count: int = 0


class Dataset(DatasetBase):
    """数据集响应模式"""
    id: int
    file_path: str
    record_count: int
    feature_count: int
    class_count: int
    user_id: int
    created_at: datetime
    
    class Config:
        orm_mode = True


class TaskStatusEnum(str, Enum):
    """任务状态枚举"""
    pending = "pending"
    running = "running"
    completed = "completed"
    failed = "failed"


class TaskBase(BaseModel):
    """检测任务基础模式"""
    algorithm_id: int
    dataset_id: int


class TaskCreate(TaskBase):
    """检测任务创建模式"""
    pass


class Task(TaskBase):
    """检测任务响应模式"""
    id: int
    status: TaskStatusEnum
    accuracy: Optional[float] = None
    recall: Optional[float] = None
    loss: Optional[float] = None
    user_id: int
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None
    created_at: datetime
    
    class Config:
        orm_mode = True


class ResultTypeEnum(str, Enum):
    """结果类型枚举"""
    traffic = "traffic"  # 流量分布
    attack = "attack"    # 攻击类型
    protocol = "protocol" # 协议类型
    flow = "flow"        # 流量波动


class ResultBase(BaseModel):
    """检测结果基础模式"""
    result_type: ResultTypeEnum
    category: str
    value: int


class ResultCreate(ResultBase):
    """检测结果创建模式"""
    task_id: int


class Result(ResultBase):
    """检测结果响应模式"""
    id: int
    task_id: int
    created_at: datetime
    
    class Config:
        orm_mode = True


class TaskDetailResponse(BaseModel):
    """任务详情响应模式"""
    task: Task
    algorithm: Algorithm
    dataset: Dataset
    results: List[Result] = []
    
    class Config:
        orm_mode = True


class TaskStatistics(BaseModel):
    """任务统计响应模式"""
    accuracy: float = 0
    recall: float = 0
    loss: float = 0
    traffic_stats: List[dict] = []  # [{"category": "HTTP", "value": 4000}, ...]
    attack_stats: List[dict] = []   # [{"category": "DDoS", "value": 300}, ...]
    protocol_stats: List[dict] = [] # [{"category": "TCP", "value": 3500}, ...]
    flow_data: List[dict] = []      # [{"timestamp": "...", "value": 1000}, ...] 